from setting import backtest
from setting import data
import polars as pl
import pickle
from setting.crate import *

# 策略名称 - 用于生成测试报告
strategy_name = "Alice3"

# 合约配置
base_data_symbol = ["SHFE_ag00_60"]  # 基准合约列表(用于时间轴)
trade_data_symbol = ["CZCE_AP00_60", "CZCE_CF00_60", "CZCE_CJ00_60"]
extra_data_symbol = []

# 加载数据
base_data, all_trade_data = data.get_data(
    base_data_symbol, 
    trade_data_symbol + extra_data_symbol, 
    "20250101",
    "20251001"
)

# 分割数据
trade_data = {k: v for k, v in all_trade_data.items() if k in trade_data_symbol}
extra_data = {k: v for k, v in all_trade_data.items() if k in extra_data_symbol}


# 回测参数配置
initial_cash = 10000 * 10000.0  # 初始资金:1000万
fee_rate = 3.0 / 10000  # 手续费率:万分之三

# 加载合约参数
with open("setting/volume_multiple.pkl", "rb") as f:
    loaded_volume = pickle.load(f)
with open("setting/margin_ratio.pkl", "rb") as f:
    loaded_margin = pickle.load(f)

deposit_rates = loaded_margin
multipliers = loaded_volume

for symbol, data in trade_data.items():
    trade_data[symbol] = data.lazy().with_columns([
        # 判断是否是周一早上10点
        ((pl.col("datetime").dt.weekday() == 1) &  # 周一 = 1
            (pl.col("datetime").dt.hour() == 10) &  # 10点
            (pl.col("datetime").dt.minute() == 0))  # 0分
        .cast(pl.Int8).alias("buy"),

        # 判断是否是周五下午2点
        ((pl.col("datetime").dt.weekday() == 5) &  # 周五 = 5
            (pl.col("datetime").dt.hour() == 14) &  # 14点
            (pl.col("datetime").dt.minute() == 0))  # 0分
        .cast(pl.Int8).alias("sell"),

        # 判断是否是周二早上10点
        ((pl.col("datetime").dt.weekday() == 2) &  # 周二 = 2
            (pl.col("datetime").dt.hour() == 10) &  # 10点
            (pl.col("datetime").dt.minute() == 0))  # 0分
        .cast(pl.Int8).alias("buy_short"),

        # 判断是否是周四早上11点
        ((pl.col("datetime").dt.weekday() == 4) &  # 周四 = 4
            (pl.col("datetime").dt.hour() == 11) &  # 11点
            (pl.col("datetime").dt.minute() == 0))  # 0分
        .cast(pl.Int8).alias("sell_short")
    ]).collect()
class Strategy(PythonStrategy):
    """
    基于时间规则的策略：周一买，周五卖，周二卖空，周四平空
    """

    def __init__(self, backtest_instance, trade_data, extra_data):
        super().__init__(backtest_instance, trade_data, extra_data)

    def on_bar(self, timestamp):
        for symbol in self.trade_data.keys():
            recent_data = self.get_recent_data(symbol, timestamp, lookback=2, require_exact_match=True)
            if recent_data is not None:
                self.calculate_time_signal(extract_symbol_base(symbol), recent_data, timestamp)

    def calculate_time_signal(self, symbol, recent_data, timestamp):
        buy = recent_data["buy"][-1]  # 周一10点买入信号
        sell = recent_data["sell"][-1]  # 周五14点卖出信号
        buy_short = recent_data["buy_short"][-1]  # 周二10点卖空信号
        sell_short = recent_data["sell_short"][-1]  # 周四11点平空信号

        # 执行交易信号
        if buy == 1:
            self.backtest_instance.buy(symbol, 2, recent_data["close"][-1])
            
        if sell == 1:
            self.backtest_instance.sell(symbol, 2, recent_data["close"][-1])
            
        if buy_short == 1:
            self.backtest_instance.buy_short(symbol, 2, recent_data["close"][-1])
            
        if sell_short == 1:
            self.backtest_instance.sell_short(symbol, 2, recent_data["close"][-1])

def main():
    """
    主函数 - 配置并运行回测
    """
    # 初始化回测引擎
    my_backtest = backtest.BackTest(initial_cash, fee_rate, deposit_rates, multipliers, trade_data)

    # 创建策略实例 - 使用子类Strategy
    strategy = Strategy(my_backtest, trade_data, extra_data)

    # 准备回测时间序列
    datetime_strings = base_data.get_column("datetime").dt.strftime("%Y-%m-%d %H:%M:%S")

    # 启动回测
    backtest.start_backtest(
        strategy,
        my_backtest,
        datetime_strings,
        base_data.get_column("timestamps"),
        strategy_name
    )

if __name__ == "__main__":
    main()